from typing import Union, Type, List, Tuple

import torch
from dynamic_network_architectures.building_blocks.residual_encoders import ResidualEncoder
from dynamic_network_architectures.building_blocks.residual import BasicBlockD, BottleneckD
from torch import nn
from torch.nn.modules.conv import _ConvNd
from torch.nn.modules.dropout import _DropoutNd

from .encoder_denoise import PlainConvEncoder
from .decoder_denoise import UNetDecoder
from dynamic_network_architectures.building_blocks.helper import convert_conv_op_to_dim

import math 

def get_timestep_embedding(timesteps, embedding_dim):
    """
    This matches the implementation in Denoising Diffusion Probabilistic Models:
    From Fairseq.
    Build sinusoidal embeddings.
    This matches the implementation in tensor2tensor, but differs slightly
    from the description in Section 3.5 of "Attention Is All You Need".
    """
    assert len(timesteps.shape) == 1

    half_dim = embedding_dim // 2
    emb = math.log(10000) / (half_dim - 1)
    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
    emb = emb.to(device=timesteps.device)
    emb = timesteps.float()[:, None] * emb[None, :]
    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
    if embedding_dim % 2 == 1:  # zero pad
        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
    return emb


def nonlinearity(x):
    # swish
    return x*torch.sigmoid(x)

class PlainConvUNet(nn.Module):
    def __init__(self,
                 input_channels: int,
                 n_stages: int,
                 features_per_stage: Union[int, List[int], Tuple[int, ...]],
                 conv_op: Type[_ConvNd],
                 kernel_sizes: Union[int, List[int], Tuple[int, ...]],
                 strides: Union[int, List[int], Tuple[int, ...]],
                 n_conv_per_stage: Union[int, List[int], Tuple[int, ...]],
                 num_classes: int,
                 n_conv_per_stage_decoder: Union[int, Tuple[int, ...], List[int]],
                 conv_bias: bool = False,
                 norm_op: Union[None, Type[nn.Module]] = None,
                 norm_op_kwargs: dict = None,
                 dropout_op: Union[None, Type[_DropoutNd]] = None,
                 dropout_op_kwargs: dict = None,
                 nonlin: Union[None, Type[torch.nn.Module]] = None,
                 nonlin_kwargs: dict = None,
                 deep_supervision: bool = False,
                 nonlin_first: bool = False
                 ):
        """
        nonlin_first: if True you get conv -> nonlin -> norm. Else it's conv -> norm -> nonlin
        """
        super().__init__()
        if isinstance(n_conv_per_stage, int):
            n_conv_per_stage = [n_conv_per_stage] * n_stages
        if isinstance(n_conv_per_stage_decoder, int):
            n_conv_per_stage_decoder = [n_conv_per_stage_decoder] * (n_stages - 1)
        assert len(n_conv_per_stage) == n_stages, "n_conv_per_stage must have as many entries as we have " \
                                                  f"resolution stages. here: {n_stages}. " \
                                                  f"n_conv_per_stage: {n_conv_per_stage}"
        assert len(n_conv_per_stage_decoder) == (n_stages - 1), "n_conv_per_stage_decoder must have one less entries " \
                                                                f"as we have resolution stages. here: {n_stages} " \
                                                                f"stages, so it should have {n_stages - 1} entries. " \
                                                                f"n_conv_per_stage_decoder: {n_conv_per_stage_decoder}"
        self.encoder = PlainConvEncoder(input_channels, n_stages, features_per_stage, conv_op, kernel_sizes, strides,
                                        n_conv_per_stage, conv_bias, norm_op, norm_op_kwargs, dropout_op,
                                        dropout_op_kwargs, nonlin, nonlin_kwargs, return_skips=True,
                                        nonlin_first=nonlin_first)
        self.decoder = UNetDecoder(self.encoder, num_classes, n_conv_per_stage_decoder, deep_supervision,
                                   nonlin_first=nonlin_first)
         # timestep embedding
        self.temb = nn.Module()
        self.temb.dense = nn.ModuleList([
            torch.nn.Linear(128,
                            input_channels),
            # torch.nn.Linear(512,
            #                 512),
        ])

    def forward(self, x, t, image=None, embeddings=None):

        temb = get_timestep_embedding(t, 128)
        temb = self.temb.dense[0](temb)
        temb = nonlinearity(temb)
        # temb = self.temb.dense[1](temb)

        # add the noisy step information
        x = x + temb[:, :, None, None, None]

        if image is not None:
            x = torch.cat([image, x], dim=1)

        x = x + temb[:, :, None, None, None]

        skips = self.encoder(x, embeddings)
        return self.decoder(skips)
    

    def compute_conv_feature_map_size(self, input_size):
        assert len(input_size) == convert_conv_op_to_dim(self.encoder.conv_op), "just give the image size without color/feature channels or " \
                                                            "batch channel. Do not give input_size=(b, c, x, y(, z)). " \
                                                            "Give input_size=(x, y(, z))!"
        return self.encoder.compute_conv_feature_map_size(input_size) + self.decoder.compute_conv_feature_map_size(input_size)


# class ResidualEncoderUNet(nn.Module):
#     def __init__(self,
#                  input_channels: int,
#                  n_stages: int,
#                  features_per_stage: Union[int, List[int], Tuple[int, ...]],
#                  conv_op: Type[_ConvNd],
#                  kernel_sizes: Union[int, List[int], Tuple[int, ...]],
#                  strides: Union[int, List[int], Tuple[int, ...]],
#                  n_blocks_per_stage: Union[int, List[int], Tuple[int, ...]],
#                  num_classes: int,
#                  n_conv_per_stage_decoder: Union[int, Tuple[int, ...], List[int]],
#                  conv_bias: bool = False,
#                  norm_op: Union[None, Type[nn.Module]] = None,
#                  norm_op_kwargs: dict = None,
#                  dropout_op: Union[None, Type[_DropoutNd]] = None,
#                  dropout_op_kwargs: dict = None,
#                  nonlin: Union[None, Type[torch.nn.Module]] = None,
#                  nonlin_kwargs: dict = None,
#                  deep_supervision: bool = False,
#                  block: Union[Type[BasicBlockD], Type[BottleneckD]] = BasicBlockD,
#                  bottleneck_channels: Union[int, List[int], Tuple[int, ...]] = None,
#                  stem_channels: int = None
#                  ):
#         super().__init__()
#         if isinstance(n_blocks_per_stage, int):
#             n_blocks_per_stage = [n_blocks_per_stage] * n_stages
#         if isinstance(n_conv_per_stage_decoder, int):
#             n_conv_per_stage_decoder = [n_conv_per_stage_decoder] * (n_stages - 1)
#         assert len(n_blocks_per_stage) == n_stages, "n_blocks_per_stage must have as many entries as we have " \
#                                                   f"resolution stages. here: {n_stages}. " \
#                                                   f"n_blocks_per_stage: {n_blocks_per_stage}"
#         assert len(n_conv_per_stage_decoder) == (n_stages - 1), "n_conv_per_stage_decoder must have one less entries " \
#                                                                 f"as we have resolution stages. here: {n_stages} " \
#                                                                 f"stages, so it should have {n_stages - 1} entries. " \
#                                                                 f"n_conv_per_stage_decoder: {n_conv_per_stage_decoder}"
#         self.encoder = ResidualEncoder(input_channels, n_stages, features_per_stage, conv_op, kernel_sizes, strides,
#                                        n_blocks_per_stage, conv_bias, norm_op, norm_op_kwargs, dropout_op,
#                                        dropout_op_kwargs, nonlin, nonlin_kwargs, block, bottleneck_channels,
#                                        return_skips=True, disable_default_stem=False, stem_channels=stem_channels)
#         self.decoder = UNetDecoder(self.encoder, num_classes, n_conv_per_stage_decoder, deep_supervision)

#     def forward(self, x, t, image=None, embeddings=None):
#         temb = get_timestep_embedding(t, 128)
#         temb = self.temb.dense[0](temb)
#         temb = nonlinearity(temb)
#         temb = self.temb.dense[1](temb)

#         if image is not None:
#             x = torch.cat([image, x], dim=1)


#         skips = self.encoder(x, temb)
#         return self.decoder(skips, temb)

#     def compute_conv_feature_map_size(self, input_size):
#         assert len(input_size) == convert_conv_op_to_dim(self.encoder.conv_op), "just give the image size without color/feature channels or " \
#                                                                                 "batch channel. Do not give input_size=(b, c, x, y(, z)). " \
#                                                                                 "Give input_size=(x, y(, z))!"
#         return self.encoder.compute_conv_feature_map_size(input_size) + self.decoder.compute_conv_feature_map_size(input_size)


if __name__ == '__main__':
    data = torch.rand((1, 4, 128, 128, 128))

    model = PlainConvUNet(4, 6, (32, 64, 125, 256, 320, 320), nn.Conv3d, 3, (1, 2, 2, 2, 2, 2), (2, 2, 2, 2, 2, 2), 4,
                                (2, 2, 2, 2, 2), False, nn.BatchNorm3d, None, None, None, nn.ReLU, deep_supervision=True)

    if False:
        import hiddenlayer as hl

        g = hl.build_graph(model, data,
                           transforms=None)
        g.save("network_architecture.pdf")
        del g

    print(model.compute_conv_feature_map_size(data.shape[2:]))

    data = torch.rand((1, 4, 512, 512))

    model = PlainConvUNet(4, 8, (32, 64, 125, 256, 512, 512, 512, 512), nn.Conv2d, 3, (1, 2, 2, 2, 2, 2, 2, 2), (2, 2, 2, 2, 2, 2, 2, 2), 4,
                                (2, 2, 2, 2, 2, 2, 2), False, nn.BatchNorm2d, None, None, None, nn.ReLU, deep_supervision=True)

    if False:
        import hiddenlayer as hl

        g = hl.build_graph(model, data,
                           transforms=None)
        g.save("network_architecture.pdf")
        del g

    print(model.compute_conv_feature_map_size(data.shape[2:]))
